import numpy as np
import os
import output_function as of
import pandas as pd
import matplotlib.pyplot as plt


# 使用示例
base_path = './data/cavity_Re_1000/'  # 根据需要修改为您的保存路径
out_iter =40000  # 假设我们要加载第300时间步的数据

# 加载数据
psi, omega, u, v, out_iter_loaded = of.load_psi_omega_uv(base_path, out_iter)
xx, yy,x,y = of.load_init_mesh(load_path=f'{base_path}initial_conditions.npz')

# 加载模拟参数
config = of.load_simulation_config(f'{base_path}simulation_config.json')
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Nx = config['Nx']
Ny = config['Ny']
Re=config['Re']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

# 打印一些变量以确认它们已被正确加载和赋值
print(f"dx: {dx}")
print(f"Re: {Re}")
print(f"dt: {dt}")
# 确认是否成功加载了正确的时间步长
if out_iter_loaded == out_iter:
    print(f"Successfully loaded data for time step: {out_iter_loaded}")
    print(omega.shape)
else:
    print("Failed to load the correct time step or file does not exist.")

# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
t_total=out_iter*dt #总时间
#绘制流函数二维颜色图。
flag_contourf=True
if flag_contourf:
    fig, ax = plt.subplots()
    # 绘制等值线图，增加levels参数来使等值线更密集  
    levels = np.linspace(psi.min(), psi.max(), 10)  # 这里设置为50个级别  
    cs = ax.contourf(xx, yy, psi, levels=levels, cmap=plt.get_cmap('Spectral').reversed())  
    plt.streamplot(xx, yy, u, v, linewidth=0.5,color='black')
    # 添加标题  
    ax.set_title('Re={},t={}'.format(Re,t_total), fontsize=16)  # 设置标题文本和字体大小  
    ax.set_xlabel('X')  # 自行替换为合适的标签
    ax.set_ylabel('Y')  # 自行替换为合适的标签
    #添加colorbar
    cbar = fig.colorbar(cs)
    cbar.set_label('psi') 
    plt.show()


v_slice = v[Ny//2, :]
# 创建图表
plt.figure()
plt.plot(x, v_slice)
plt.xlabel('x')
plt.ylabel('v-midplane')  # 这里假设y轴是某种值，您可以根据实际情况修改
plt.title('Re={},t={}'.format(Re,t_total))
plt.show()